<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Build tfjs model</title>

    <script src="../../assets/libs/tf.min.js"></script>
    <script src="mnist_data.js"></script>

</head>
<body>

<script>
	// Build basic model structure
	const input = tf.input({shape: [28, 28, 1]});
	const conv1 = tf.layers.conv2d({
		kernelSize: 5,
		filters: 6,
		strides: 1,
		activation: 'relu',
		kernelInitializer: 'VarianceScaling',
		name: 'MyConv2D_1'
	});
	const maxPool1 = tf.layers.maxPooling2d({
		poolSize: [2, 2],
		strides: [2, 2],
		name: 'MyMaxPooling_1'
	});
	const conv2 = tf.layers.conv2d({
		kernelSize: 5,
		filters: 16,
		strides: 1,
		activation: 'relu',
		kernelInitializer: 'VarianceScaling',
		name: 'MyConv2D_2'
	});
	const maxPool2 = tf.layers.maxPooling2d({
		poolSize: [2, 2],
		strides: [2, 2],
		name: 'MyMaxPooling_2'
	});
	const flatten = tf.layers.flatten();
	const dense1 = tf.layers.dense({
		units: 120,
		kernelInitializer: 'VarianceScaling',
		activation: 'relu',
		name: 'MyDense_1'
	});
	const dense2 = tf.layers.dense({
		units: 84,
		kernelInitializer: 'VarianceScaling',
		activation: 'relu',
		name: 'MyDense_2'
	});
	const dense3 = tf.layers.dense({
		units: 10,
		kernelInitializer: 'VarianceScaling',
		activation: 'softmax',
		name: 'MySoftMax'
	});
	// Apply input for each layer and collect the output
	const conv1Output = conv1.apply(input);
	const maxPool1Output = maxPool1.apply(conv1Output);
	const conv2Output = conv2.apply(maxPool1Output);
	const maxPool2Output = maxPool2.apply(conv2Output);
	const flattenOutput = flatten.apply(maxPool2Output);
	const dense1Output = dense1.apply(flattenOutput);
	const dense2Output = dense2.apply(dense1Output);
	const dense3Output = dense3.apply(dense2Output);

	const model = tf.model({
		inputs: input,
		outputs: dense3Output
	});
	// Set training config and compile
	const LEARNING_RATE = 0.0001;
	const optimizer = tf.train.adam(LEARNING_RATE);
	model.compile({
		optimizer: optimizer,
		loss: 'categoricalCrossentropy',
		metrics: ['accuracy'],
	});
	// Load training data
	let data;
	async function load() {
		data = new MnistData();
		await data.load();
	}
	async function train() {
		const BATCH_SIZE = 50;
		const TRAIN_BATCHES = 200;
		const TEST_BATCH_SIZE = 1000;
		const TEST_ITERATION_FREQUENCY = 100;
		for (let i = 0; i < TRAIN_BATCHES; i++) {
			const batch = data.nextTrainBatch(BATCH_SIZE);
			let testBatch;
			let validationData;
			if (i % TEST_ITERATION_FREQUENCY === 0) {
				testBatch = data.nextTestBatch(TEST_BATCH_SIZE);
				validationData = [
					testBatch.xs.reshape([TEST_BATCH_SIZE, 28, 28, 1]), testBatch.labels
				];
			}
			const history = await model.fit(
				batch.xs.reshape([BATCH_SIZE, 28, 28, 1]),
				batch.labels,
				{
					batchSize: BATCH_SIZE,
					validationData,
					epochs: 1
				});
			if (i % TEST_ITERATION_FREQUENCY === 0) {
				const loss = history.history.loss[0];
				const accuracy = history.history.acc[0];
				console.log(accuracy);
			}
		}
	}
	async function saveModel() {
		await model.save("downloads://mnist");
	}
	async function mnist() {
		await load();
		await train();
		await saveModel();
	}
	mnist();
</script>

</body>
</html>